import json,_thread
from django_redis import get_redis_connection

from dvadmin.knowledge.models import KnowledgeRecord,KnowledgeRecordSegment

class RedisUtil():

    def __init__(self):
        self.conn = get_redis_connection("default")
        
        self.channel = 'knowledge'        # 定义频道
        print("初始化Redis连接")
        sub_obj = self.conn.pubsub()          
        sub_obj.subscribe(self.channel)  
        sub_obj.parse_response()          # 丢掉第一个消息，这个是订阅成功的消息
        _thread.start_new_thread(self.dualEmbedding, (sub_obj,)) 
    
    def dualEmbedding(self,sub_obj):
        while True:                                    # 一直接收消息
            try :
                msg = sub_obj.parse_response()             # 得到消息
                print("Received Msg: %s" % msg)            # 打印消息
                rid = msg[2]
                knowledgeRecordSegments = KnowledgeRecordSegment.objects.filter(rid=rid)
                
                chunkList=[]
                data_vector=[]
                
                for knowledgeRecordSegment in knowledgeRecordSegments:
                    kid = knowledgeRecordSegment.get("kid")
                    chunkList.append(knowledgeRecordSegment.get("content"))
                    data_vector.append({
                        "id": knowledgeRecordSegment.get("id"),
                        "kid": kid,
                        "file_id": knowledgeRecordSegment.get("rid"),
                        "file_valid": 1,
                        "chunk_valid": 1,   
                        "context": knowledgeRecordSegment.get("content")
                    })
                
                # 向量化
                embeddings = self.embedding.embed_documents(chunkList)
                
                row = 0
                for embedding in embeddings:            
                    data_vector[row].set({
                        "embedding":embedding
                    })
                
                tablename = "knowledge" + str(kid)
                # 保存到向量数据库
                self.milvus.upset(tablename,data_vector)
                
                KnowledgeRecord.objects.get(id=rid).update(status=1)
            except Exception as e:
                print(e)
                KnowledgeRecord.objects.get(id=rid).update(status=0)